package org.offsync.async

import java.net.Socket
import java.net.SocketAddress
import java.nio.channels.{Selector => JavaSelector}
import java.nio.channels.{SelectionKey => JavaSelectionKey}
import java.nio.channels.{SocketChannel => JavaSocketChannel}
import java.nio.ByteBuffer
import java.net.InetSocketAddress
import org.offsync.rich._
import org.offsync.util.Window
import org.offsync.util.DisposedException
import java.nio.channels.SelectionKey

case class AsyncSocket(
    override val channel: JavaSocketChannel)
    extends AsyncSelectable(channel) {
  private val connectInterest = new AsyncInterest(this, SelectionKey.OP_CONNECT)
  private val readInterest = new AsyncInterest(this, SelectionKey.OP_READ)
  private val writeInterest = new AsyncInterest(this, SelectionKey.OP_WRITE)
  
  override val ops = List(connectInterest, readInterest, writeInterest)
  
  def streamVia(selector: AsyncSelector): AsyncStream = new AsyncStream() {
    override def writeSome(window: Window[Byte])(task: (() => Int) => Unit): Unit = {
      AsyncSocket.this.writeSome(selector, window)(task)
    }
  
    override def writeFully(window: Window[Byte])(task: (() => Int) => Unit): Unit = {
      AsyncSocket.this.writeFully(selector, window)(task)
    }
  
    override def readSome(window: Window[Byte])(task: (() => Int) => Unit): Unit = {
      AsyncSocket.this.readSome(selector, window)(task)
    }
  
    override def readFully(window: Window[Byte])(task: (() => Int) => Unit): Unit = {
      AsyncSocket.this.readFully(selector, window)(task)
    }
  }
  
  private def fastWrite(window: Window[Byte], fastWriteCompletion: () => Unit)(task: (() => Int) => Unit): Unit = {
    val bytesWritten = try {
      fastWriteCompletion()

      channel.write(window.toByteBuffer())
    } catch {
      case e => task { () => throw e.wrap() }
      return
    }
    
    task { () => bytesWritten }
  }
  
  def writeSome(selector: AsyncSelector, window: Window[Byte])(task: (() => Int) => Unit): Unit = {
    fastWrite(window, { () => }) { fastWriteCompletion =>
      val bytesWritten = try {
        fastWriteCompletion()
      } catch {
        case e => task { () => throw e.wrap() }
        return
      }
      
      if (bytesWritten == 0 && window.length != 0) {
        writeInterest.onReady(selector) { readyCompletion =>
          try {
            readyCompletion()
          } catch {
            case e => {
              task { () => throw e.wrap() }
              return
            }
          }
          writeSome(selector, window)(task)
        }
      } else {
        task { () => bytesWritten }
      }
    }
  }
  
  def writeFully(selector: AsyncSelector, window: Window[Byte])(task: (() => Int) => Unit): Unit = {
    fastWrite(window, {() => }) { writeFullyCompletion =>
      val bytesWritten = try {
        writeFullyCompletion()
      } catch {
        case e => {
          task { () => throw e.wrap() }
          return
        }
      }
      
      if (bytesWritten < window.length && bytesWritten >= 0) {
        writeInterest.onReady(selector) { onReadyCompletion =>
          try {
            onReadyCompletion()
          } catch {
            case e => {
              task { () => throw e.wrap() }
              return
            }
          }
          writeFully(selector, window.from(bytesWritten))(task)
        }
      } else {
        task { () => bytesWritten }
      }
    }
  }
  
  private def fastRead(window: Window[Byte], fastReadCompletion: () => Unit)(task: (() => Int) => Unit): Unit = {
    val bytesRead = try {
      fastReadCompletion();

      channel.read(window.toByteBuffer());
    } catch {
      case e => {
        task { () => throw e.wrap() }
        return
      }
    }
    
    task { () => bytesRead }
  }
  
  def readSome(selector: AsyncSelector, window: Window[Byte])(task: (() => Int) => Unit) {
    fastRead(window, { () => }) { readSomeCompletion =>
      val bytesRead = try {
        readSomeCompletion()
      } catch {
        case e => {
          task { () => throw e.wrap() }
          return
        }
      }
      
      if (bytesRead == 0 && window.length != 0) {
        readInterest.onReady(selector) { onReadyCompletion =>
          try {
            onReadyCompletion()
          } catch {
            case e => {
              task { () => throw e.wrap() }
              return
            }
          }
          readSome(selector, window)(task)
        }
      } else {
        task { () => bytesRead }
      }
    }
  }
  
  def readFully(selector: AsyncSelector, window: Window[Byte])(task: (() => Int) => Unit): Unit = {
    fastRead(window, {() => }) { readCompletion =>
      val bytesRead = readCompletion();
      
      if (bytesRead < window.length && bytesRead >= 0) {
        readInterest.onReady(selector) { onReadyCompletion =>
          try {
            onReadyCompletion()
          } catch {
            case e => {
              task { () => throw e.wrap() }
              return
            }
          }
          readFully(selector, window.from(bytesRead))(task)
        }
      } else {
        task { () => bytesRead }
      }
    }
  }
  
  private def fastConnect(fastConnectCompletion: () => Unit)(task: (() => Boolean) => Unit): Unit = {
    val connected = try {
      fastConnectCompletion()

      channel.finishConnect()
    } catch {
      case e => {
        task { () => throw e.wrap() }
        return
      }
    }
    
    task { () => connected }
  }
  
  private def finishConnect(selector: AsyncSelector)(task: (() => Boolean) => Unit): Unit = {
    fastConnect({() => }) { connectCompletion =>
      val connected = connectCompletion()
      
      if (!connected) {
        connectInterest.onReady(selector) { onReadyCompletion =>
          try {
            onReadyCompletion()
          } catch {
            case e => {
              task { () => throw e.wrap() }
              return
            }
          }
          
          finishConnect(selector)(task);
        }
      } else {
        task { () => connected }
      }
    }
  }
  
  def connect(selector: AsyncSelector, remote: SocketAddress, tcpNoDelay: Boolean)(task: (() => Boolean) => Unit) {
    try {
      channel.connect(remote);
      channel.socket().setTcpNoDelay(tcpNoDelay);
      finishConnect(selector)(task)
    } catch {
      case e => task { () => throw e.wrap() }
      return;
    }
  }
  
  def connect(selector: AsyncSelector, host: String, port: Int, tcpNoDelay: Boolean)(task: (() => Boolean) => Unit) {
    connect(selector, new InetSocketAddress(host, port), tcpNoDelay)(task)
  }
  
  protected override def onDispose(): Unit = {
    try {
      channel.close();
    } catch {
      case e => e.printStackTrace(System.err)
    } finally {
      super.onDispose()
    }
  }
}

object AsyncSocket {
  def open(): AsyncSocket = {
    val channel = JavaSocketChannel.open()
    channel.configureBlocking(false)
    return AsyncSocket(channel)
  }
  
  def open(remote: SocketAddress): AsyncSocket = {
    val channel = JavaSocketChannel.open(remote)
    channel.configureBlocking(false)
    AsyncSocket(channel)
  }

  def open(host: String, port: Int): AsyncSocket = open(new InetSocketAddress(host, port))
}
